from pathlib import Path

from langchain_core.messages import BaseMessage
from langchain_core.prompts import ChatPromptTemplate
_prompt: dict[str, str] = {}


def get_prompt(name: str, **kwargs) -> list[BaseMessage]:
    global _prompt
    if name not in _prompt:
        _prompt[name] = _load_prompt(name)
    return ChatPromptTemplate(
        [("system", _prompt[name])],
        template_format="jinja2"
    ).invoke(kwargs).to_messages()


def _load_prompt(name: str) -> str:
    file_path = Path(__file__).parent / f"{name}.md"
    with open(file_path, "r", encoding="utf-8") as f:
        prompt = f.read()
    return prompt


if __name__ == "__main__":
    if 'controller' not in _prompt:
        _prompt['controller'] = _load_prompt('controller')
    p = ChatPromptTemplate([("system", _prompt['controller'])]).invoke({"language": 'cn'})
    print(p.to_messages())
